type AnyObject = Record<PropertyKey, any>;
import XRequest from '../x-request';
import type { XRequestParams } from '../x-request';
import type { SSEOutput, XStreamOptions } from '../x-stream';
interface RequestFnInfo<Message> extends AnyObject {
  messages?: Message[];
  message?: Message;
}

export type RequestFn<Input, Output> = (
  info: XRequestParams & Input,
  callbacks: {
    onUpdate: (chunk: Output) => void;
    onSuccess: (chunks: Output[]) => void;
    onError: (error: Error) => void;
    onStream?: (abortController: AbortController) => void;
  },
  transformStream?: XStreamOptions<Output>['transformStream']
) => void;

export interface XAgentConfigPreset {
  baseURL: string;
  key: string;
  model: string;
  dangerouslyApiKey: string;
}
export interface XAgentConfigCustom<Input, Output> {
  request?: RequestFn<Input, Output>;
}

export type XAgentConfig<Input, Output> = Partial<XAgentConfigPreset> &
  XAgentConfigCustom<Input, Output>;

let uuid = 0;

/** This is a wrap class to avoid developer can get too much on origin object */
export class XAgent<
  Message = string,
  Input = RequestFnInfo<Message>,
  Output = SSEOutput
> {
  config: XAgentConfig<Input, Output>;

  private requestingMap: Record<number, boolean> = {};

  constructor(config: XAgentConfig<Input, Output>) {
    this.config = config;
    if (!this.config.request) {
      const xRequest = XRequest({
        baseURL: this.config.baseURL!,
        model: this.config.model,
        dangerouslyApiKey: this.config.dangerouslyApiKey,
      }).create;
      this.config.request = (
        info: XRequestParams & Input,
        callbacks: {
          onUpdate: (chunk: Output) => void;
          onSuccess: (chunks: Output[]) => void;
          onError: (error: Error) => void;
          onStream?: (abortController: AbortController) => void;
        },
        transformStream?
      ) => {
        xRequest(info, callbacks, transformStream);
      };
    }
  }

  private finishRequest(id: number) {
    delete this.requestingMap[id];
  }

  public request: RequestFn<Input, Output> = (
    info,
    callbacks,
    transformStream?
  ) => {
    const { request } = this.config;
    const { onUpdate, onSuccess, onError, onStream } = callbacks;

    const id = uuid;
    uuid += 1;
    this.requestingMap[id] = true;

    request?.(
      info,
      {
        onStream: (abortController) => {
          if (this.requestingMap[id]) {
            onStream?.(abortController);
          }
        },
        // Status should be unique.
        // One get success or error should not get more message
        onUpdate: (chunk) => {
          if (this.requestingMap[id]) {
            onUpdate(chunk);
          }
        },
        onSuccess: (chunks) => {
          if (this.requestingMap[id]) {
            onSuccess(chunks);
            this.finishRequest(id);
          }
        },
        onError: (error) => {
          if (this.requestingMap[id]) {
            onError(error);
            this.finishRequest(id);
          }
        },
      },
      transformStream
    );
  };

  public isRequesting() {
    return Object.keys(this.requestingMap).length > 0;
  }
}
